import cv2
import time
import subprocess as sp
import multiprocessing
import psutil
import ffmpeg


class stream_pusher(object):
    def __init__(self, rtmp_url=None, raw_frame_q=None, fps=25, width=1920, height=1080):  # 类实例化的时候传入rtmp地址和帧传入队列
        self.rtmp_url = rtmp_url
        self.raw_frame_q = raw_frame_q
        self.fps = fps  # 设置帧速率
        # 设置分辨率
        self.width = int(width)  # 宽
        self.height = int(height)  # 高

        # 设置FFmpeg命令文本
        self.command = ['ffmpeg',
                        '-y',
                        '-c', 'copy',                   # copy图像的质量会更好
                        # '-f', 'rawvideo',
                        # '-vcodec','rawvideo',
                        '-pix_fmt', 'bgr24',
                        '-max_delay', str(100),
                         '-s', "{}x{}".format(self.width, self.height),
                        '-r', str(fps),
                        '-i', '-',
                        '-c:v', 'libx264',
                        '-pix_fmt', 'yuv420p',
                        '-preset', 'ultrafast',
                        '-f', 'flv',
                        '-g', '5',
                        '-b', '7000000',
                        self.rtmp_url]

    # 对获取的帧做一些画面处理的方法，返回完成处理的帧。
    def __frame_handle__(self, raw_frame, text, shape1, shape2):
        # 帧用cv2进行一些处理，比如写上文本，画矩形等
        return (raw_frame)

    # 向服务器推送
    def push_frame(self):
        p = psutil.Process()
        p.cpu_affinity([0, 1, 2, 3])
        # 配置向os传递命令的管道
        p = sp.Popen(self.command, stdin=sp.PIPE)

        while True:
            if not self.raw_frame_q.empty():  # 如果输入管道不为空
                # 把帧和相关信息从输入队列中取出
                raw_frame, text, shape1, shape2 = self.raw_frame_q.get()
                # 对获取的帧进行画面处理
                frame = self.__frame_handle__(raw_frame, text, shape1, shape2)
                # 把内容放入管道，放入后有os自己去执行
                p.stdin.write(frame.tostring())
            else:
                time.sleep(0.01)

    # 启动运行
    def run(self):
        # 定义一个子进程
        push_frame_p = multiprocessing.Process(target=self.push_frame, args=())
        push_frame_p.daemon = True  # 把子进程设置为daemon方式
        push_frame_p.start()  # 运行子进程